/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.mapreduce;

import java.io.IOException;
import java.util.List;

import jtokyotyrant.JTT;
import jtokyotyrant.RDB;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * implementation of db that uses tokyo tyrant as the backend
 * @author nakada
 *
 */
public class TTDB implements DB {
  private static Logger logger = LoggerFactory.getLogger(TTDB.class);
  private final RDB db;
  private final DBI dbImpl;

  private interface DBI {
    public void putlist(byte[][] kvs) throws IOException;
    public List<byte[]> range(byte[] prefix, byte[] sentinel, int count,
                              boolean binc) throws IOException;
    public List<byte[]> ranged(byte[] prefix, byte[] sentinel, int count,
                               boolean binc) throws IOException;
    public int rangeout(byte[] prefix, byte[] sentinel, int count,
                        boolean binc) throws IOException;
    public List<byte[]> getlist(byte[][] keys) throws IOException;
    public void outlist(byte[][] keys) throws IOException;
  }

  public TTDB(IOConfiguration ioConf, String host, int port) {
    this.db = JTT.rdbnew();
    try {
      this.db.open(host, port);
    }
    catch (OutOfMemoryError e) {
      logger.error("open {}", db.errmsg(db.ecode()));
      throw e;
    }
    this.dbImpl = ioConf.db_use_vanilla ?
      new DBI() {               // vanilla Tokyo Tyrant
        @Override
        public void putlist(byte[][] kvs) throws IOException{
          try {
            db.misc("putlist", kvs);
          }
          catch (OutOfMemoryError e) {
            logger.error("putlist {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public List<byte[]> range(byte[] prefix, byte[] sentinel, int count,
                                  boolean binc) throws IOException {
          logger.error("TTDB#range is not implemented for vanilla Tokyo Tyrant/Cabinet");
          throw new IOException
            (new SssException
             ("safety misc.range is not supported by vanilla Tokyo Tyrant/Cabinet"));
        }
        @Override
        public List<byte[]> ranged(byte[] prefix, byte[] sentinel, int count,
                                   boolean binc) throws IOException {
          logger.error("TTDB#ranged is not implemented for vanilla Tokyo Tyrant/Cabinet");
          throw new IOException
            (new SssException
             ("safety misc.ranged is not supported by vanilla Tokyo Tyrant/Cabinet"));
        }
        @Override
        public int rangeout(byte[] prefix, byte[] sentinel, int count,
                            boolean binc) throws IOException{
          logger.error("TTDB#rangeout is not implemented for vanilla Tokyo Tyrant/Cabinet");
          throw new IOException
            (new SssException
             ("misc.rangeout is not supported by vanilla Tokyo Tyrant/Cabinet"));
        }
        @Override
        public List<byte[]> getlist(byte[][] keys) throws IOException {
          try {
            return db.misc("getlist", keys);
          }
          catch (OutOfMemoryError e) {
            logger.error("getlist {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public void outlist(byte[][] keys) throws IOException {
          try {
            db.misc("outlist", keys);
          }
          catch (OutOfMemoryError e) {
            logger.error("outlist {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
      } :
      new DBI() {               // extended Tokyo Tyrant
        @Override
        public void putlist(byte[][] kvs) throws IOException{
          try {
            db.misc("putlist_atomic2", kvs);
          }
          catch (OutOfMemoryError e) {
            logger.error("putlist_atomic2 {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public List<byte[]> range(byte[] prefix, byte[] sentinel, int limit,
                                  boolean binc) throws IOException {
          try {
            return db.misc("range_size_atomic",
                           prefix,
                           Integer.toString(limit).getBytes("UTF-8"),
                           sentinel,
                           Integer.toString(binc ? 1 : 0).getBytes("UTF-8"));
          }
          catch (OutOfMemoryError e) {
            logger.error("range_size_atomic {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public List<byte[]> ranged(byte[] prefix, byte[] sentinel, int limit,
                                   boolean binc) throws IOException {
          try {
            return db.misc("range_size_atomic_d",
                           prefix,
                           Integer.toString(limit).getBytes("UTF-8"),
                           sentinel,
                           Integer.toString(binc ? 1 : 0).getBytes("UTF-8"));
          }
          catch (OutOfMemoryError e) {
            logger.error("range_size_atomic_d {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public int rangeout(byte[] prefix, byte[] sentinel, int count,
                            boolean binc) throws IOException{
          try {
            List<byte[]> r = db.misc
              ("rangeout_atomic", prefix,
               Integer.toString(count).getBytes("UTF-8"), sentinel,
               Integer.toString(binc ? 1 : 0).getBytes("UTF-8"));
            return r.size() == 0 ? 0 : Integer.parseInt(new String(r.get(0)));
          }
          catch (NumberFormatException e) {
            throw new IOException("Invalid result", e);
          }
          catch (OutOfMemoryError e) {
            logger.error("rangeout_atomic{}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public List<byte[]> getlist(byte[][] keys) throws IOException {
          try {
            return db.misc("getlist_atomic", keys);
          }
          catch (OutOfMemoryError e) {
            logger.error("getlist_atomic{}", db.errmsg(db.ecode()));
            throw e;
          }
        }
        @Override
        public void outlist(byte[][] keys) throws IOException {
          try {
            db.misc("outlist_atomic", keys);
          }
          catch (OutOfMemoryError e) {
            logger.error("outlist_atomic {}", db.errmsg(db.ecode()));
            throw e;
          }
        }
      };
  }
  /**
   * Bulk put operation
   */
  @Override
  public void putlist(byte[][] kvs) throws IOException{
    dbImpl.putlist(kvs);
  }
  /**
   * Bulk get operation
   */
  @Override
  public List<byte[]> range(byte[] prefix, byte[] sentinel, int count,
                            boolean binc) throws IOException {
    return dbImpl.range(prefix, sentinel, count, binc);
  }
  /**
   * Bulk get operation (destructive)
   */
  @Override
  public List<byte[]> ranged(byte[] prefix, byte[] sentinel, int count,
                             boolean binc) throws IOException {
    return dbImpl.ranged(prefix, sentinel, count, binc);
  }
  /**
   * Bulk delete operation
   */
  @Override
  public int rangeout(byte[] prefix, byte[] sentinel, int count, boolean binc)
    throws IOException{
    return dbImpl.rangeout(prefix, sentinel, count, binc);
  }
  @Override
  public List<byte[]> getlist(byte[][] keys) throws IOException {
    return dbImpl.getlist(keys);
  }
  @Override
  public void outlist(byte[][] keys) throws IOException {
    dbImpl.outlist(keys);
  }

  @Override
  public List<byte[]> fwmkeys(byte[] prefix, int max) throws IOException {
    try {
      return db.fwmkeys(prefix, max);
    }
    catch (OutOfMemoryError e) {
      logger.error("fwmkeys{}", db.errmsg(db.ecode()));
      throw e;
    }
  }

  @Override
  public byte[] get(byte[] key) throws IOException {
    try {
      return db.get(key);
    }
    catch (OutOfMemoryError e) {
      logger.error("get{}", db.errmsg(db.ecode()));
      throw e;
    }
  }

  @Override
  public void close() throws IOException {
    try {
      db.close();
    }
    catch (OutOfMemoryError e) {
      logger.error("close{}", db.errmsg(db.ecode()));
      throw e;
    }
  }
}
